"""
一个使用 LangChain 构建的智能体，结合搜索引擎功能
"""

import os
import sys
from typing import Dict, List, Optional
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.llms.tongyi import Tongyi
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tracers import ConsoleCallbackHandler
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_community.tools import TavilySearchResults
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.tools import Tool

# 设置标准输入输出的编码为 UTF-8
if sys.platform == 'win32':
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
    sys.stdin = codecs.getreader('utf-8')(sys.stdin.buffer)

# 加载环境变量
load_dotenv()

class SearchAgent:
    def __init__(self, debug: bool = False):
        """
        初始化智能体
        
        Args:
            debug: 是否启用调试模式，启用后会在控制台显示详细的调用信息
        """
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("请设置 DASHSCOPE_API_KEY 环境变量")
        
        # 设置回调处理器
        self.callbacks = [ConsoleCallbackHandler()] if debug else None
        
        # 初始化通义大模型
        self.llm = Tongyi(
            model_name="qwen-plus",  # 可选: qwen-turbo, qwen-plus, qwen-max
            temperature=0.7,
            api_key=self.api_key,
            callbacks=self.callbacks
        )

        # 创建搜索引擎工具
        self.search = TavilySearchResults()
        self.tools = [
            Tool(
                name="search",
                func=self.search.run,
                description="用于搜索互联网上的信息。输入应该是一个搜索查询。"
            )
        ]

        # 创建提示词模板
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个智能助手，可以使用搜索引擎来回答问题。
            当你需要查找信息时，使用搜索工具。
            请用简洁明了的方式回答问题。
            如果搜索结果不够准确，请说明情况。
            
            可用工具:
            {tools}
            
            工具名称: {tool_names}
            
            历史对话:
            {history}
            
            当前问题: {input}
            
            思考过程:
            {agent_scratchpad}
            
            请按照以下格式回答：
            Thought: 思考过程
            Action: 要使用的工具名称
            Action Input: 工具的输入
            Observation: 工具的返回结果
            Thought: 根据结果继续思考
            Final Answer: 最终答案""")
        ])

        # 创建智能体
        self.agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=self.prompt
        )

        # 创建执行器
        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            verbose=debug,
            handle_parsing_errors=True
        )

    def run(self, user_input: str, history: List[Dict] = None) -> str:
        """
        运行智能体
        
        Args:
            user_input: 用户输入的问题
            history: 对话历史，格式为 [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
            
        Returns:
            str: 智能体的回答
            
        Raises:
            Exception: 当处理过程中发生错误时
        """
        try:
            # 转换历史记录格式
            messages = []
            if history:
                for msg in history:
                    if msg["role"] == "user":
                        messages.append(HumanMessage(content=msg["content"]))
                    elif msg["role"] == "assistant":
                        messages.append(AIMessage(content=msg["content"]))
            
            response = self.agent_executor.invoke(
                {
                    "input": user_input,
                    "history": messages
                },
                config={
                    "callbacks": self.callbacks,
                    "metadata": {"user_input": user_input}
                }
            )
            return response["output"]
        except Exception as e:
            return f"发生错误: {str(e)}"

    def get_chain(self):
        """获取智能体链，用于 LangServe 部署"""
        return self.agent_executor

def create_agent(debug: bool = False) -> SearchAgent:
    """
    创建智能体实例
    
    Args:
        debug: 是否启用调试模式
        
    Returns:
        SearchAgent: 智能体实例
    """
    return SearchAgent(debug=debug)

def main():
    """主函数"""
    # 确保终端编码为 UTF-8
    if sys.platform == 'win32':
        os.system('chcp 65001')
    
    agent = create_agent(debug=True)
    
    print("欢迎使用智能搜索助手！输入 'quit' 退出。")
    print("当前使用的是 qwen-plus 模型，您可以在代码中修改为 qwen-turbo 或 qwen-max")
    
    history = []
    while True:
        try:
            user_input = input("\n请输入您的问题: ")
            if user_input.lower() == 'quit':
                break
                
            response = agent.run(user_input, history)
            print(f"\nAI 助手: {response}")
            
            # 更新历史记录
            history.append({"role": "user", "content": user_input})
            history.append({"role": "assistant", "content": response})
        except UnicodeDecodeError:
            print("输入编码错误，请重试")
        except Exception as e:
            print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 